// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/memory_reclaimer.h>

#include <pollux/exec/driver.h>
#include <pollux/exec/task.h>

namespace kumo::pollux::exec {
    MemoryReclaimer::MemoryReclaimer(int32_t priority)
        : memory::MemoryReclaimer(priority) {
    }

    std::unique_ptr<memory::MemoryReclaimer> MemoryReclaimer::create(
        int32_t priority) {
        return std::unique_ptr<memory::MemoryReclaimer>(
            new MemoryReclaimer(priority));
    }

    void MemoryReclaimer::enterArbitration() {
        DriverThreadContext *driverThreadCtx = driverThreadContext();
        if (MELON_UNLIKELY(driverThreadCtx == nullptr)) {
            // Skips the driver suspension handling if this memory arbitration
            // request is not issued from a driver thread.
            return;
        }

        Driver *const driver = driverThreadCtx->driverCtx()->driver;
        if (driver->task()->enterSuspended(driver->state()) != StopReason::kNone) {
            // There is no need for arbitration if the associated task has already
            // terminated.
            POLLUX_FAIL("Terminate detected when entering suspension");
        }
    }

    void MemoryReclaimer::leaveArbitration() noexcept {
        DriverThreadContext *driverThreadCtx = driverThreadContext();
        if (MELON_UNLIKELY(driverThreadCtx == nullptr)) {
            // Skips the driver suspension handling if this memory arbitration
            // request is not issued from a driver thread.
            return;
        }
        Driver *const driver = driverThreadCtx->driverCtx()->driver;
        driver->task()->leaveSuspended(driver->state());
    }

    void MemoryReclaimer::abort(
        memory::MemoryPool *pool,
        const std::exception_ptr &error) {
        if (pool->kind() == memory::MemoryPool::Kind::kLeaf) {
            return;
        }
        pool->visitChildren([&](memory::MemoryPool *child) {
            auto *reclaimer = child->reclaimer();
            if (reclaimer != nullptr) {
                reclaimer->abort(child, error);
            }
            return true;
        });
    }

    /*static*/
    std::unique_ptr<memory::MemoryReclaimer>
    ParallelMemoryReclaimer::create(melon::Executor *executor, int32_t priority) {
        return std::unique_ptr<memory::MemoryReclaimer>(
            new ParallelMemoryReclaimer(executor, priority));
    }

    ParallelMemoryReclaimer::ParallelMemoryReclaimer(
        melon::Executor *executor,
        int32_t priority)
        : MemoryReclaimer(priority), executor_(executor) {
    }

    uint64_t ParallelMemoryReclaimer::reclaim(
        memory::MemoryPool *pool,
        uint64_t targetBytes,
        uint64_t maxWaitMs,
        Stats &stats) {
        if (executor_ == nullptr) {
            return memory::MemoryReclaimer::reclaim(
                pool, targetBytes, maxWaitMs, stats);
        }

        // Sort candidates based on priority.
        struct Candidate {
            std::shared_ptr<memory::MemoryPool> pool;
            int64_t reclaimableBytes;
        };
        std::vector<Candidate> candidates; {
            std::shared_lock guard{pool->poolMutex_};
            candidates.reserve(pool->children_.size());
            for (auto &entry: pool->children_) {
                auto child = entry.second.lock();
                if (child != nullptr) {
                    const auto reclaimableBytesOpt = child->reclaimableBytes();
                    if (!reclaimableBytesOpt.has_value()) {
                        continue;
                    }
                    candidates.push_back(Candidate{
                        std::move(child),
                        static_cast<int64_t>(reclaimableBytesOpt.value())
                    });
                }
            }
        }

        struct ReclaimResult {
            const uint64_t reclaimedBytes{0};
            const Stats stats;
            const std::exception_ptr error{nullptr};

            explicit ReclaimResult(std::exception_ptr _error)
                : reclaimedBytes(0), error(std::move(_error)) {
            }

            ReclaimResult(uint64_t _reclaimedBytes, Stats &&_stats)
                : reclaimedBytes(_reclaimedBytes),
                  stats(std::move(_stats)),
                  error(nullptr) {
            }
        };

        std::vector<std::shared_ptr<AsyncSource<ReclaimResult> > > reclaimTasks;
        for (const auto &candidate: candidates) {
            if (candidate.reclaimableBytes == 0) {
                continue;
            }
            reclaimTasks.push_back(memory::createAsyncMemoryReclaimTask<ReclaimResult>(
                [&, reclaimPool = candidate.pool]() {
                    try {
                        Stats reclaimStats;
                        const auto bytes =
                                reclaimPool->reclaim(targetBytes, maxWaitMs, reclaimStats);
                        return std::make_unique<ReclaimResult>(
                            bytes, std::move(reclaimStats));
                    } catch (const std::exception &e) {
                        POLLUX_MEM_LOG(ERROR) << "Reclaim from memory pool " << pool->name()
                                 << " failed: " << e.what();
                        // The exception is captured and thrown by the caller.
                        return std::make_unique<ReclaimResult>(std::current_exception());
                    }
                }));
            if (reclaimTasks.size() > 1) {
                executor_->add([source = reclaimTasks.back()]() { source->prepare(); });
            }
        }

        auto syncGuard = melon::makeGuard([&]() {
            for (auto &reclaimTask: reclaimTasks) {
                // We consume the result for the pending tasks. This is a cleanup in the
                // guard and must not throw. The first error is already captured before
                // this runs.
                try {
                    reclaimTask->move();
                } catch (const std::exception &) {
                }
            }
        });

        uint64_t reclaimedBytes{0};
        for (auto &reclaimTask: reclaimTasks) {
            const auto result = reclaimTask->move();
            if (result->error) {
                std::rethrow_exception(result->error);
            }
            stats += result->stats;
            reclaimedBytes += result->reclaimedBytes;
        }
        return reclaimedBytes;
    }

    void memoryArbitrationStateCheck(memory::MemoryPool &pool) {
        const auto *driverThreadCtx = driverThreadContext();
        if (driverThreadCtx != nullptr) {
            Driver *driver = driverThreadCtx->driverCtx()->driver;
            if (!driver->state().suspended()) {
                POLLUX_FAIL(
                    "Driver thread is not suspended under memory arbitration processing: {}, request memory pool: {}",
                    driver->toString(),
                    pool.name());
            }
        }
    }
} // namespace kumo::pollux::exec
